package main

import (
    "errors"
    "fmt"
)

// 定义堆结构体
type Heap[T int] struct {
    arr []T
}

// 创建堆
func NewHeap[T int](arr []T) *Heap[T] {
    // 拷贝数组
    brr := make([]T, len(arr))
	copy(brr, arr)
    return &Heap[T]{arr: brr}
}

// 向下调整
func (h *Heap[T]) downAdjust(parent int) {
    left := 2*parent + 1 // 左孩子在数组里面的下标
    if left >= len(h.arr) { // 左孩子不存在了，直接返回
        return
    }
    // 从父子三人中找最小者
    minIndex := parent
    minValue := h.arr[parent]

    if h.arr[left] < minValue {
        minIndex = left
        minValue = h.arr[left]
    }

    right := left + 1
    if right < len(h.arr) {
        if h.arr[right] < minValue {
            minIndex = right
        }
    }
    // 最小值跟父节点交换
    if minIndex != parent {
        h.arr[minIndex], h.arr[parent] = h.arr[parent], h.arr[minIndex]
        h.downAdjust(minIndex) // 递归调用,每当有新的元素插入，就要对他的三角区域进行下沉调整
    }
}

// 向上调整
func (h *Heap[T]) upAdjust(child int) {
    if child == 0 {
        return
    }
    parent := (child - 1) / 2
    // 如果父节点大于子节点，则交换
    if h.arr[parent] > h.arr[child] {
        h.arr[parent], h.arr[child] = h.arr[child], h.arr[parent]
        h.upAdjust(parent) // 递归调用,每当有新的元素插入，就要对他的三角区域进行上浮调整
    }
}

// 构建堆
func (h *Heap[T]) BuildHeap() {
    n := len(h.arr)
    if n <= 1 {
        return
    }
    lastIndex := (n / 2) - 1 // 最后一个非叶子节点的下标
    for i := lastIndex; i >= 0; i-- {
        h.downAdjust(i) // 下沉调整
    }
}

// 插入一个元素
func (h *Heap[T]) Insert(value T) {
    h.arr = append(h.arr, value) // 追加一个元素
    h.upAdjust(len(h.arr) - 1)   // 上浮调整
}

// 删除并返回顶堆元素
func (h *Heap[T]) Pop() (T, error) {
    if len(h.arr) == 0 {
        var v T
        return v, errors.New("heap is empty")
    }
    root := h.arr[0]
    // 直接用最后一个元素替换顶堆元素
    h.arr[0] = h.arr[len(h.arr)-1]
    h.arr = h.arr[:len(h.arr)-1] // 删除最后一个元素
    // 向下调整
    h.downAdjust(0)
    return root, nil
}

// 取得顶堆元素
func (h *Heap[T]) Peek() (T, error) {
    if len(h.arr) == 0 {
        var v T
        return v, errors.New("heap is empty")
    }
    return h.arr[0], nil
}

// 获取堆的大小
func (h *Heap[T]) Size() int {
    return len(h.arr)
}

// 获取堆里面所有元素
func (h *Heap[T]) GetAll() []T {
    return h.arr
}

// 替换顶堆元素
func (h *Heap[T]) ReplaceTop(ele T) {
    if len(h.arr) == 0 {
        return
    }
    h.arr[0] = ele
    h.downAdjust(0)
}

func main() {
    // 示例代码
    heap := NewHeap[int]([]int{3, 2, 1, 5, 6, 4})
    heap.BuildHeap()
    fmt.Println(heap.Peek()) // 输出: 1
    fmt.Println(heap.Pop())  // 输出: 1, nil
    fmt.Println(heap.Peek()) // 输出: 2
}